/**
 * 
 */
package edu.umd.clip.lm.model;

import java.io.*;
import java.util.*;

import org.w3c.dom.*;

import edu.umd.clip.lm.model.decoding.Decoder;
import edu.umd.clip.lm.util.IO;
import edu.umd.clip.lm.util.XMLUtils;
import edu.umd.clip.lm.util.XMLizable;
import edu.umd.clip.lm.util.tree.BinaryTree;
import edu.umd.clip.lm.util.tree.BinaryTreeIterator;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ForestModel extends XMLizable {
	private static final long serialVersionUID = 1L;

	public static final String XML_ELEMENT_NAME = "forest-model";
	private static final String XML_TREES_NAME = "trees";
	private static final String XML_MODEL_NAME = "model";
	private static final String XML_BACKOFF_ATTR_NAME = "use-grandparent-for-backoff";
	private static final String XML_NAME_ATTR = "name";
	
	private ArrayList<LanguageModel> models;
	private HashMap<LanguageModel, Integer> modelNumbers;
	private boolean backoffLeavesUseGrandparent = true;
	private byte hiddenOrder;
	private byte overtOrder;
	private String name;
	private Decoder decoder;
	
	private DecodingRuntime decodingRuntime;
	
	/**
	 * @param xmlnode
	 */
	public ForestModel(Element xmlnode) {
		assert(xmlnode.getTagName().equals(XML_ELEMENT_NAME));
		name = xmlnode.getAttribute(XML_NAME_ATTR);
		
		String attr = xmlnode.getAttribute(XML_BACKOFF_ATTR_NAME);
		backoffLeavesUseGrandparent = Boolean.parseBoolean(attr);
		
		Element elem = XMLUtils.getFirstElementByTagName(xmlnode, XML_TREES_NAME);
		NodeList nodes = XMLUtils.getChildrenElementsByTagName(elem, XML_MODEL_NAME);
		models = new ArrayList<LanguageModel>(nodes.getLength());
		modelNumbers = new HashMap<LanguageModel,Integer>(models.size());
		
		Experiment exp = Experiment.getInstance();
		
		hiddenOrder = 0;
		overtOrder = 0;
		for(int i=0; i<nodes.getLength(); ++i) {
			String modelName = ((Element) nodes.item(i)).getTextContent();
			LanguageModel lm = exp.getLM(modelName);
			addLM(lm);
		}
	}

	public void addLM(LanguageModel lm) {
		models.add(lm);
		modelNumbers.put(lm, modelNumbers.size());
		
		if (lm.getHiddenOrder() > hiddenOrder) hiddenOrder = lm.getHiddenOrder();
		if (lm.getOvertOrder() > overtOrder) overtOrder = lm.getOvertOrder();
	}
	/**
	 * 
	 */
	public ForestModel(String name) {
		this.name = name;
		this.models = new ArrayList<LanguageModel>();
		this.modelNumbers = new HashMap<LanguageModel, Integer>();
		this.hiddenOrder = 0;
		this.overtOrder = 0;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.util.XMLizable#createXML(org.w3c.dom.Document)
	 */
	@Override
	public Element createXML(Document doc) {
		Element element = doc.createElement(XML_ELEMENT_NAME);
		
		element.setAttribute(XML_NAME_ATTR, name);
		if (backoffLeavesUseGrandparent) {
			element.setAttribute(XML_BACKOFF_ATTR_NAME, Boolean.TRUE.toString());
		}
		
		Element trees = doc.createElement(XML_TREES_NAME);
		for(LanguageModel lm : models) {
			Element lmElem = doc.createElement(XML_MODEL_NAME);
			lmElem.setTextContent(lm.getId());
			trees.appendChild(lmElem);
		}
		element.appendChild(trees);
		
		return element;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.util.XMLizable#getElementName()
	 */
	@Override
	public String getElementName() {
		return XML_ELEMENT_NAME;
	}

	public boolean isBackoffLeavesUseGrandparent() {
		return backoffLeavesUseGrandparent;
	}

	public void setBackoffLeavesUseGrandparent(boolean backoffLeavesUseGrandparent) {
		this.backoffLeavesUseGrandparent = backoffLeavesUseGrandparent;
	}

	public ArrayList<LanguageModel> getModels() {
		return models;
	}

	public byte getHiddenOrder() {
		return hiddenOrder;
	}

	public byte getOvertOrder() {
		return overtOrder;
	}

	public String getName() {
		return name;
	}

	public byte getOrder() {
		return hiddenOrder > overtOrder ? hiddenOrder : overtOrder;
	}

	public Decoder getDecoder() {
		if (decoder == null) {
			synchronized(this) {
				if (decoder == null) {
					decoder = new Decoder(this);
				}
			}
		}
		return decoder;
	}
	
	public static class DecodingRuntime {
		private final double weights[][];
		//private final BinaryTree<HistoryTreePayload> nodes[][];
		
		/**
		 * @param weights
		 */
		private DecodingRuntime(double[][] weights) {
			this.weights = weights;
		}
		
		public final double[] getWeights(int lmNum) {
			return weights[lmNum];
		}
		
		public void setWeights(int lmNum, double lmWeights[]) {
			this.weights[lmNum] = lmWeights;
		}
	}
	
	private void loadDecodingRuntime() {
		double weights[][];
		try {
			weights = loadInterpolationWeights();
			assert(weights.length == models.size());
		} catch(IOException e) {
			weights = new double[models.size()][];
			for(int lmNum=0; lmNum<models.size(); ++lmNum) {
				LanguageModel lm = models.get(lmNum);
				weights[lmNum] = new double[lm.getNodes().length];
			}
		}
		decodingRuntime = new DecodingRuntime(weights);
	}
	
	public String getInterpolationWeightsFilename() {
		return getName() + "-weights.ser.gz";
	}
	
	private double[][] loadInterpolationWeights() throws IOException {
		ObjectInputStream ois = new ObjectInputStream(IO.getInputStream(getInterpolationWeightsFilename()));
		try {
			return (double[][]) ois.readObject();
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
			System.exit(1);
		}
		ois.close();
		return null;
	}
	
	public DecodingRuntime getDecodingRuntime() {
		if (decodingRuntime == null) {
			synchronized(this) {
				if (decodingRuntime == null) {
					loadDecodingRuntime();
				}
			}
		}
		return decodingRuntime;
	}
	
	public void saveInterpolationWeights() throws IOException {
		getDecodingRuntime();
		
		ObjectOutputStream oos = new ObjectOutputStream(IO.getOutputStream(getInterpolationWeightsFilename()));
		oos.writeObject(decodingRuntime.weights);
		oos.close();
	}

	public int getModelNumber(LanguageModel lm) {
		return modelNumbers.get(lm).intValue();
	}
}
